#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
模型下载脚本
用于下载BERT模型到本地目录
"""

import os
from transformers import AutoTokenizer, AutoModelForTokenClassification
import shutil

def download_model(model_name: str = "chinese-bert-wwm-ext",
                   save_path: str = "/home/aresen/1project/2python/hub/models/chinese-bert-wwm-ext"):
                   # save_path: str = "./models/bert-base-chinese"):

    """
    下载BERT模型到指定路径

    Args:
        model_name: 要下载的模型名称
        save_path: 保存路径
    """
    print(f"开始下载模型 {model_name} 到 {save_path}")

    # 创建保存目录
    os.makedirs(save_path, exist_ok=True)

    try:
        # 下载分词器
        print("下载分词器...")
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        tokenizer.save_pretrained(save_path)
        print("分词器下载完成")

        # 下载模型
        print("下载模型...")
        model = AutoModelForTokenClassification.from_pretrained(model_name)
        model.save_pretrained(save_path)
        print("模型下载完成")

        print(f"模型已成功下载到: {save_path}")
        print("目录结构:")
        for root, dirs, files in os.walk(save_path):
            level = root.replace(save_path, '').count(os.sep)
            indent = ' ' * 2 * level
            print(f"{indent}{os.path.basename(root)}/")
            subindent = ' ' * 2 * (level + 1)
            for file in files:
                print(f"{subindent}{file}")

    except Exception as e:
        print(f"下载模型时出错: {e}")
        # 清理不完整的下载
        if os.path.exists(save_path):
            shutil.rmtree(save_path)
        raise

if __name__ == "__main__":
    download_model()
